Skip to content

[NKI-352] Make sure BIR AP lowers fully in lower AP pass #572

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
ppotapov-aws merged 1 commit into from
Jan 15, 2026
Merged

[NKI-352] Make sure BIR AP lowers fully in lower AP pass #572

ppotapov-aws merged 1 commit into from
Jan 15, 2026

Conversation

@ppotapov-aws
Copy link
Collaborator

@ppotapov-aws ppotapov-aws commented Jan 15, 2026 •
edited
Loading

We need to lower BIR AP properly since vector offset now is an Access it has to be lowered as well.
Sample kernel:

def kernel():
  a = nl.ndarray((32, 32), nl.int8)
  b = nl.ndarray((32, 32), nl.int8)
  c = nl.ndarray((32, 32), nl.int8)
  nisa.dma_copy(a, b.ap([[32, 1], [1,1]], vector_offset=c.ap([[2,2]])))

Output copy component for b in KLR

src := KLR.Core.TensorRef.abstract
         (KLR.Core.Access.birPattern
           { tensor := { name := "b.3",
                         dtype := KLR.Core.Dtype.int8,
                         shape := { parDim := 32, freeDims := [32] },
                         address := { name := "b.3",
                                      memory := KLR.Core.Memory.sbuf,
                                      parSize := 32,
                                      freeSize := 32,
                                      parOffset := none,
                                      freeOffset := none,
                                      isShared := false },
                         freeElements := 32,
                         parWF := _,
                         freeWF := _,
                         addressRotation := false },
             offset := 0,
             pattern := [{ step := 32, num := 1, offset := 0 },
                         { step := 1, num := 1, offset := 0 }],
             scalarOffset := none,
             vectorOffset := some (KLR.Core.Access.birPattern
                               { tensor := { name := "c.4",
                                             dtype := KLR.Core.Dtype.int8,
                                             shape := { parDim := 32,
                                                        freeDims := [32] },
                                             address := { name := "c.4",
                                                          memory := KLR.Core.Memory.sbuf,
                                                          parSize := 32,
                                                          freeSize := 32,
                                                          parOffset := none,
                                                          freeOffset := none,
                                                          isShared := false },
                                             freeElements := 32,
                                             parWF := _,
                                             freeWF := _,
                                             addressRotation := false },
                                 offset := 66,
                                 pattern := [{ step := 32, num := 1, offset := 2 },
                                             { step := 1, num := 1, offset := 2 }],
                                 scalarOffset := none,
                                 vectorOffset := none,
                                 indirectDim := 0,
                                 dtypeOverride := none }),
                                              indirectDim := 0,
                                              dtypeOverride := none }),
                                 compute_op := KLR.Core.DgeComputeOp.none,
                                 oobMode := KLR.Core.DmaBounds.error,
                                 dgeMode := 0,
                                 uniqueIndices := false,
                                 engine := KLR.Core.Engine.unassigned })

@yongweiy
Copy link
Collaborator

yongweiy commented Jan 15, 2026

Can you explain what's c[[2,2]]?

@ppotapov-aws
Copy link
Collaborator Author

ppotapov-aws commented Jan 15, 2026 •
edited
Loading

Can you explain what's c[[2,2]]?

Just something to make it look like a vector offset. Ah, I forgot to add .ap([[2,2]]). Should be copy pasting and not typing by hand. Fixed the description

@ppotapov-aws ppotapov-aws merged commit b348060 into main Jan 15, 2026
6 checks passed
@ppotapov-aws ppotapov-aws deleted the ppotapov/lower_bir_ap branch January 15, 2026 18:20
@ppotapov-aws ppotapov-aws restored the ppotapov/lower_bir_ap branch January 20, 2026 18:26
@ppotapov-aws ppotapov-aws mentioned this pull request Jan 20, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yongweiy yongweiy yongweiy approved these changes

@govereau govereau Awaiting requested review from govereau govereau is a code owner

@seanmcl seanmcl Awaiting requested review from seanmcl seanmcl is a code owner

@jtristan jtristan Awaiting requested review from jtristan jtristan is a code owner

@joonwonc joonwonc Awaiting requested review from joonwonc joonwonc is a code owner

@waahm7 waahm7 Awaiting requested review from waahm7 waahm7 is a code owner

@jiancheng-aws jiancheng-aws Awaiting requested review from jiancheng-aws jiancheng-aws is a code owner

@kerrijoe-aws kerrijoe-aws Awaiting requested review from kerrijoe-aws kerrijoe-aws is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ppotapov-aws @yongweiy